package com.sau.base;

import dev.langchain4j.chain.ConversationalChain;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;

import java.math.BigDecimal;
import java.sql.*;
import java.util.HashMap;
import java.util.Map;
import java.util.Scanner;

public class StockAnalysisWithModel {

    // 数据库连接信息
    private static final String DB_URL = "jdbc:mysql://localhost:3306/java?useSSL=false&serverTimezone=UTC";
    private static final String DB_USER = "root";
    private static final String DB_PASSWORD = "root";

    // 初始化Ollama聊天模型
    private ChatLanguageModel model;
    private ConversationalChain chain;

    public StockAnalysisWithModel() {
        // 创建 Ollama 聊天模型实例
        this.model = OllamaChatModel.builder()
                .baseUrl("http://127.0.0.1:11434")
                .modelName("qwen2:0.5b")
                .build();

        // 创建对话内存，用于存储对话上下文
        ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(10);

        // 创建 ConversationalChain，将模型和对话内存绑定
        this.chain = ConversationalChain.builder()
                .chatLanguageModel(model)
                .chatMemory(chatMemory)
                .build();
    }

    // 定义提示词模板
    private static final String TEMPLATE = "你是一个专业的金融助手。请根据以下股票信息提供详细的投资分析：\n\n" +
            "{{stockInfo}}\n\n" +
            "请从以下角度分析：\n" +
            "1. 当前股价和交易所\n" +
            "2. 最新分红记录\n" +
            "3. 近期股价趋势\n" +
            "4. 投资建议\n\n" +
            "请确保分析清晰、简洁。";

    // 主方法：执行股票分析
    public void analyzeStock() {
        Scanner scanner = new Scanner(System.in);
        System.out.print("请输入股票代码（例如：000001）：");
        String stockCode = scanner.nextLine().trim();

        // 查询数据库信息
        Map<String, String> stockInfo = queryStockInfo(stockCode);

        // 如果没有查询到信息
        if (stockInfo.isEmpty()) {
            System.out.println("未找到该股票的相关信息。");
            return;
        }

        // 构建提示词
        String prompt = buildPrompt(stockInfo);
        System.out.println("当前提示词为："+prompt);

        // 调用大模型生成回复
        String response = chain.execute(prompt);

        // 输出模型回复
        System.out.println("模型分析结果：");
        System.out.println(response);
    }

    // 查询股票信息
    private Map<String, String> queryStockInfo(String stockCode) {
        Map<String, String> stockInfo = new HashMap<>();

        try (Connection conn = DriverManager.getConnection(DB_URL, DB_USER, DB_PASSWORD)) {
            // 查询股票基本信息
            String basicSql = "SELECT stock_nm, exchange FROM m_stocks WHERE stock_code = ?";
            try (PreparedStatement pstmt = conn.prepareStatement(basicSql)) {
                pstmt.setString(1, stockCode);
                try (ResultSet rs = pstmt.executeQuery()) {
                    if (rs.next()) {
                        stockInfo.put("股票名称", rs.getString("stock_nm"));
                        stockInfo.put("交易所", rs.getString("exchange"));
                    }
                }
            }

            // 查询最新分红记录
            String dividendSql = "SELECT dividend FROM divided_record WHERE stock_code = ? ORDER BY announcement_date DESC LIMIT 1";
            try (PreparedStatement pstmt = conn.prepareStatement(dividendSql)) {
                pstmt.setString(1, stockCode);
                try (ResultSet rs = pstmt.executeQuery()) {
                    if (rs.next()) {
                        stockInfo.put("最新分红", rs.getString("dividend"));
                    }
                }
            }

            // 查询近期股价趋势
            String historySql = "SELECT close_price, trade_date FROM stocks_history WHERE stock_code = ? ORDER BY trade_date DESC LIMIT 5";
            try (PreparedStatement pstmt = conn.prepareStatement(historySql)) {
                pstmt.setString(1, stockCode);
                try (ResultSet rs = pstmt.executeQuery()) {
                    StringBuilder trend = new StringBuilder();
                    while (rs.next()) {
                        trend.append(rs.getString("trade_date"))
                                .append(": ")
                                .append(rs.getString("close_price"))
                                .append("\n");
                    }
                    stockInfo.put("近期股价趋势", trend.toString());
                }
            }

        } catch (SQLException e) {
            e.printStackTrace();
        }

        return stockInfo;
    }

    // 构建提示词
    private String buildPrompt(Map<String, String> stockInfo) {
        StringBuilder promptBuilder = new StringBuilder();
        promptBuilder.append("股票代码: 000001").append("\n");
        promptBuilder.append("股票名称: ").append(stockInfo.get("股票名称")).append("\n");
        promptBuilder.append("交易所: ").append(stockInfo.get("交易所")).append("\n");
        promptBuilder.append("最新分红: ").append(stockInfo.get("最新分红")).append("\n");
        promptBuilder.append("近期股价趋势:\n").append(stockInfo.get("近期股价趋势")).append("\n");

        return TEMPLATE.replace("{{stockInfo}}", promptBuilder.toString());
    }

    // 主程序入口
    public static void main(String[] args) {
        StockAnalysisWithModel analyzer = new StockAnalysisWithModel();
        analyzer.analyzeStock();
    }
}
